import os
import torch
import cv2
import numpy as np
import torch.nn.functional as F
from model import model_unet_res50
from tqdm import tqdm


model = model_unet_res50.cuda()
model.load_state_dict(torch.load('../model_save/new/unet_res50_baseline/best.pth'))
save_path = '../submit/result'


os.makedirs(save_path,exist_ok=True)
test_data_path = '../suichang_round1_test_partA_210120\suichang_round1_test_partA_210120'


for name in tqdm(os.listdir(test_data_path)):
    image_dir = os.path.join(test_data_path,name)
    image = cv2.imread(image_dir)
    image = torch.from_numpy(np.transpose(image,(2,0,1))).float().cuda()[np.newaxis,...]
    pred = model(image)
    pred = torch.argmax(pred, dim=1).cpu().numpy()[0]
    ## 注意将0-9转成1-10
    pred = pred + 1
    pred = np.stack((pred,pred,pred),axis=2)
    pred = pred.astype(np.uint8)
    cv2.imwrite(os.path.join(save_path,name.replace('tif','png')),pred)
